from .swin_unet import get_swin_unet
from .unet import UNet

def get_model(args):
    if args.model =='swinunet':
        if isinstance(args.train_crop_size,list) or isinstance(args.train_crop_size,tuple):
            image_size=args.train_crop_size[0]
        else:
            image_size=args.train_crop_size
        model = get_swin_unet(image_size=image_size,num_classes=args.num_classes,in_channels=args.in_channels)
    elif args.model == 'unet':
        model = UNet(in_channels=args.in_channels,num_classes=args.num_classes)
    else:
        raise NotImplementedError

    return model


def build_model(args):
    if args.model =='swinunet':
        if isinstance(args.train_crop_size,list) or isinstance(args.train_crop_size,tuple):
            image_size=args.train_crop_size[0]
        else:
            image_size=args.train_crop_size
        model = get_swin_unet(image_size=image_size,num_classes=args.num_classes,in_channels=args.in_channels)
    elif args.model == 'unet':
        model = UNet(in_channels=args.in_channels,num_classes=args.num_classes)
    else:
        raise NotImplementedError

    return model